import cv2
import numpy as np

"""
shape: 高 宽 (通道数)
size:  宽 高
"""


def show(img, name=""):
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


img_l = cv2.imread("rsc/left.jpg")
img_r = cv2.imread("rsc/right.jpg")

ratio = 0.6
# 统一大小
width = int(img_l.shape[0] * ratio)
height = int(img_l.shape[1] * ratio)
img_l = cv2.resize(img_l, (width, height), interpolation=cv2.INTER_AREA)
img_r = cv2.resize(img_r, (width, height), interpolation=cv2.INTER_AREA)

# 获得特征点和描述子
sift = cv2.SIFT.create()
kp_l, desc_l = sift.detectAndCompute(img_l, None)
kp_r, desc_r = sift.detectAndCompute(img_r, None)

# 建立暴力匹配器
matcher = cv2.BFMatcher()

# 暴力匹配 每个匹配两个候选
matches = matcher.knnMatch(desc_l, desc_r, 2)

goods = []
for m, n in matches:
    if m.distance < 0.75 * n.distance:
        # 最短的两个距离差别比较大，那就保留这个短的匹配
        goods.append(m)

img_match = cv2.drawMatches(img_l, kp_l, img_r, kp_r, goods, None)
show(img_match, "matches")

# 拼接图片
if len(goods) > 4:
    # 获取单应性矩阵
    # 需要float32的ndarray
    pts_l = np.float32([kp_l[m.queryIdx].pt for m in goods])
    pts_r = np.float32([kp_r[m.trainIdx].pt for m in goods])
    H, mask = cv2.findHomography(pts_r, pts_l, cv2.RANSAC, 2)
    # 根据变换扭曲照片
    warp = cv2.warpPerspective(img_r, H, (img_l.shape[1] + img_r.shape[1], img_r.shape[0]))
    show(warp, "warp")
    # 拼接
    stitch = warp.copy()
    # print(img_l.shape, img_r.shape, stitch.shape)
    stitch[:img_l.shape[0], :img_l.shape[1]] = img_l
    show(stitch, "stitch")
    cv2.imwrite("rsc/lr_stitch.jpg", stitch)
